import torch
import torch.nn as nn
import numpy as np
from mayavi import mlab
from mayavi.mlab import *
import prob
import math
from NN import cdinn2

############################################
# set default data type to double
############################################
torch.set_default_dtype(torch.float64)
torch.set_default_tensor_type(torch.DoubleTensor)


###########################################
# plot function for 3-d systems
# support cube, sphere and cylinder
###########################################


# code for plotting cube: from
# https://stackoverflow.com/questions/26888098/draw-cubes-with-colour-intensity-with-python
def mlab_plt_cube(xmin, xmax, ymin, ymax, zmin, zmax, c_color):
	faces = cube_faces( trans_x(xmin), trans_x(xmax), trans_y(ymin), trans_y(ymax), trans_z(zmin), trans_z(zmax) )
	# faces = cube_faces( xmin, xmax, ymin, ymax, zmin, zmax )
	for grid in faces:
		x,y,z = grid
		mlab.mesh(x,y,z,color=c_color,opacity=0.3)


def cube_faces(xmin, xmax, ymin, ymax, zmin, zmax):
	faces = []

	x,y = np.mgrid[xmin:xmax:3j,ymin:ymax:3j]
	z = np.ones(y.shape)*zmin
	faces.append((x,y,z))

	x,y = np.mgrid[xmin:xmax:3j,ymin:ymax:3j]
	z = np.ones(y.shape)*zmax
	faces.append((x,y,z))

	x,z = np.mgrid[xmin:xmax:3j,zmin:zmax:3j]
	y = np.ones(z.shape)*ymin
	faces.append((x,y,z))

	x,z = np.mgrid[xmin:xmax:3j,zmin:zmax:3j]
	y = np.ones(z.shape)*ymax
	faces.append((x,y,z))

	y,z = np.mgrid[ymin:ymax:3j,zmin:zmax:3j]
	x = np.ones(z.shape)*xmin
	faces.append((x,y,z))

	y,z = np.mgrid[ymin:ymax:3j,zmin:zmax:3j]
	x = np.ones(z.shape)*xmax
	faces.append((x,y,z))

	return faces
def trans_x(origin_x):
	k = (prob.PLOT_LEN_B[0] - 1) / (prob.domain_max[0] - prob.domain_min[0])
	c = 1 - k * prob.domain_min[0]
	return k * origin_x + c

def trans_y(origin_y):
	k = (prob.PLOT_LEN_B[1] - 1) / (prob.domain_max[1] - prob.domain_min[1])
	c = 1 - k * prob.domain_min[1]
	return k * origin_y + c

def trans_z(origin_z):
	k = (prob.PLOT_LEN_B[2] - 1) / (prob.domain_max[2] - prob.domain_min[2])
	c = 1 - k * prob.domain_min[2]
	return k * origin_z + c


# plot sphere
def mlab_plt_sphere(center_x, center_y, center_z, s_rad, s_color):
	x, y, z = np.ogrid[ (prob.domain_min[0]) : (prob.domain_max[0]) : complex(0, prob.PLOT_LEN_B[0]), \
							(prob.domain_min[1]) : (prob.domain_max[1]) : complex(0, prob.PLOT_LEN_B[1]), \
								(prob.domain_min[2]) : (prob.domain_max[2]) : complex(0, prob.PLOT_LEN_B[2]) ]

	sphere_scalar = (x - center_x) * (x - center_x) + (y - center_y) * (y - center_y) + (z - center_z) * (z - center_z)
	sphere = mlab.contour3d(sphere_scalar, contours = [s_rad * s_rad], color = s_color, opacity = 0.3)


# plot sphere
def mlab_plt_cylinder(center_x, center_y, s_rad, s_color):
	x, y, z = np.ogrid[ (prob.domain_min[0]) : (prob.domain_max[0]) : complex(0, prob.PLOT_LEN_B[0]), \
							(prob.domain_min[1]) : (prob.domain_max[1]) : complex(0, prob.PLOT_LEN_B[1]), \
								(prob.domain_min[2]) : (prob.domain_max[2]) : complex(0, prob.PLOT_LEN_B[2]) ]

	cylinder_scalar = (x - center_x) * (x - center_x) + (y - center_y) * (y - center_y) - 0 * z
	cylinder = mlab.contour3d(cylinder_scalar, contours = [s_rad * s_rad], color = s_color, opacity = 0.3)


# plot initial
def mlab_plt_init():
	if prob.init_shape == 1: # cube
		mlab_plt_cube(prob.init_min[0], prob.init_max[0], prob.init_min[1], prob.init_max[1], prob.init_min[2], prob.init_max[2], (0, 1, 0) )
	else: # sphere
		mlab_plt_sphere( (prob.init_min[0] + prob.init_max[0]) / 2, (prob.init_min[1] + prob.init_max[1]) / 2, \
							(prob.init_min[2] + prob.init_max[2]) / 2, (prob.init_max[0] - prob.init_min[0]) / 2, (0, 1, 0) )


# plot unsafe
def mlab_plt_unsafe():
	# if len(prob.SUB_UNSAFE) == 0:

	if prob.unsafe_shape == 1: # cube
		mlab_plt_cube(prob.unsafe_min[0], prob.unsafe_max[0], prob.unsafe_min[1], prob.unsafe_max[1], prob.unsafe_min[2], prob.unsafe_max[2], (1, 0, 0) )
	elif prob.unsafe_shape == 2: # sphere
		mlab_plt_sphere( (prob.unsafe_min[0] + prob.unsafe_max[0]) / 2, (prob.unsafe_min[1] + prob.unsafe_max[1]) / 2, \
							(prob.unsafe_min[2] + prob.unsafe_max[2]) / 2, (prob.unsafe_max[0] - prob.unsafe_min[0]) / 2, (1, 0, 0) )
	elif prob.unsafe_shape == 3: # cylinder
		mlab_plt_cylinder( (prob.unsafe_min[0] + prob.unsafe_max[0]) / 2, (prob.unsafe_min[1] + prob.unsafe_max[1]) / 2, \
			(prob.unsafe_max[0] - prob.unsafe_min[0]) / 2, (1, 0, 0) )
	else:
		x, y, z = np.ogrid[ (prob.domain_min[0]) : (prob.domain_max[0]) : complex(0, prob.PLOT_LEN_B[0]), \
						(prob.domain_min[1]) : (prob.domain_max[1]) : complex(0, prob.PLOT_LEN_B[1]), \
							(prob.domain_min[2]) : (prob.domain_max[2]) : complex(0, prob.PLOT_LEN_B[2]) ]
		hyperplane_scalar = x + y + z
		hyperplane = mlab.contour3d(hyperplane_scalar, contours = [1], color = (1, 0, 0), opacity = 0.3)

	# else:
	# 	for i in range(len(prob.SUB_UNSAFE)):
	# 		curr_shape = prob.SUB_UNSAFE_SHAPE[i]
	# 		curr_range = prob.SUB_UNSAFE[i]
	# 		if curr_shape == 1: # cube
	# 			mlab_plt_cube(curr_range[0][0], curr_range[0][1], curr_range[1][0], curr_range[1][1], curr_range[2][0], curr_range[2][1], (1, 0, 0) )
	# 		elif curr_shape == 2: # sphere
	# 			mlab_plt_sphere( (curr_range[0][0] + curr_range[0][1]) / 2, (curr_range[1][0] + curr_range[1][1]) / 2, \
	# 								(curr_range[2][0] + curr_range[2][1]) / 2, (curr_range[0][1] - curr_range[0][0]) / 2, (1, 0, 0) )
	# 		else : # cylinder
	# 			mlab_plt_cylinder( (curr_range[0][0] + curr_range[0][1]) / 2, (curr_range[1][0] + curr_range[1][1]) / 2, \
	# 								 (curr_range[0][1] - curr_range[0][0]) / 2, (1, 0, 0) )


# generating plot data for nn
def gen_plot_data():
	sample_x = torch.linspace(prob.domain_min[0], prob.domain_max[0], int(prob.PLOT_LEN_B[0]))
	sample_y = torch.linspace(prob.domain_min[1], prob.domain_max[1], int(prob.PLOT_LEN_B[1]))
	sample_z = torch.linspace(prob.domain_min[2], prob.domain_max[2], int(prob.PLOT_LEN_B[2]))
	grid_xyz = torch.meshgrid([sample_x, sample_y, sample_z])
	flatten_xyz = [torch.flatten(grid_xyz[i]) for i in range(len(grid_xyz))]
	plot_input = torch.stack(flatten_xyz, 1)
	return plot_input

def gen_plot_data_init():
	sample_x = torch.linspace(prob.init_min[0], prob.init_max[0], 2)
	sample_y = torch.linspace(prob.init_min[1], prob.init_max[1], 2)
	sample_z = torch.linspace(prob.init_min[2], prob.init_max[2], 2)
	grid_xyz = torch.meshgrid([sample_x, sample_y, sample_z])
	flatten_xyz = [torch.flatten(grid_xyz[i]) for i in range(len(grid_xyz))]
	plot_input = torch.stack(flatten_xyz, 1)
	return plot_input

# plot barrier
def mlab_plt_barrier(model):
	# generating nn_output for plotting
	plot_input = gen_plot_data()
	# print("plot_input:", plot_input)
	nn_output = model(plot_input)
	# print(nn_output)
	# print("nn_output[:, 0]:")
	# print(nn_output[:, 0])
	# print("nn_output:", nn_output)
	plot_output = (nn_output[:, 0]).reshape(prob.PLOT_LEN_B[0], prob.PLOT_LEN_B[1], prob.PLOT_LEN_B[2])
	# print("plot_output:", plot_output)
	barrier_plot = mlab.contour3d(plot_output.detach().numpy(), contours=[0], color=(1, 1, 0), opacity=0.3) # yellow
	# src = mlab.pipeline.scalar_field(plot_output.detach().numpy())
	# barrier_plot = mlab.pipeline.iso_surface(src, contours=[0], color=(1, 1, 0), opacity=1)


def use_model(x,y,z):
	# return math.tanh(math.tanh(y + math.tanh(-0.48 + (-y / math.sin(y + z + math.cos(x + 2 * y))))))  # zhrR1
	return (-0.54543733999999999 + math.sin(math.sin(((x - math.tanh((0.82548980000000005 + y))) / (1.1922542933451701 + 1.33552912280763 * math.exp(( - y + z))))))) # C9

def mlab_plt_barrier_withoutModel():
	# generating nn_output for plotting
	# plot_input = gen_plot_data()
	# print("plot_input:", plot_input)

	x = np.linspace(prob.domain_min[0], prob.domain_max[0], int(prob.PLOT_LEN_B[0]))
	y = np.linspace(prob.domain_min[1], prob.domain_max[1], int(prob.PLOT_LEN_B[1]))
	z = np.linspace(prob.domain_min[2], prob.domain_max[2], int(prob.PLOT_LEN_B[2]))
	xx, yy, zz = np.meshgrid(x, y, z)
	s = np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T
	nn_output = []
	for i in range((int)(s.size / 3)):
		nn_output.append(use_model(s[i][0], s[i][1], s[i][2]))
	# nn_output = use_model(s[:, 0], s[:, 1], s[:, 2])
	# nn_output = use_model(xx, yy, zz)
	# plot_v = nn_output.detach().numpy()
	nn_output = np.array(nn_output)
	nn_output = torch.tensor(nn_output, dtype=torch.float64).to("cpu")

	print(nn_output)

	# nn_output = model(plot_input)
	# print("nn_output:", nn_output)
	plot_output = nn_output.reshape(prob.PLOT_LEN_B[0], prob.PLOT_LEN_B[1], prob.PLOT_LEN_B[2])
	# print("plot_output:", plot_output)
	barrier_plot = mlab.contour3d(plot_output.detach().numpy(), contours=[0], color=(1, 1, 0), opacity=0.3) # yellow
	# src = mlab.pipeline.scalar_field(plot_output.detach().numpy())
	# barrier_plot = mlab.pipeline.iso_surface(src, contours=[0], color=(1, 1, 0), opacity=1)


# plot vector field
def mlab_plt_vector():

	vector_input = gen_plot_data()

	vector_field = prob.vector_field(vector_input)

	# attention! for plot-3d to work, prob.PLOT_LEN_V should be equal to prob.PLOT_LEN_B
	u = (vector_field[:, 0]).reshape(prob.PLOT_LEN_B[0], prob.PLOT_LEN_B[1], prob.PLOT_LEN_B[2])
	v = (vector_field[:, 1]).reshape(prob.PLOT_LEN_B[0], prob.PLOT_LEN_B[1], prob.PLOT_LEN_B[2])
	w = (vector_field[:, 2]).reshape(prob.PLOT_LEN_B[0], prob.PLOT_LEN_B[1], prob.PLOT_LEN_B[2])
	mlab.quiver3d(u.numpy(), v.numpy(), w.numpy(), color=(0, 0, 1), mask_points=10, scale_factor=1)


def mlab_plt_discrete_trajectory():
	vector_input = gen_plot_data_init()
	# print("vector_input:", vector_input)
	if prob.init_shape == 2:
		p0 = (prob.init_min[0] + prob.init_max[0]) / 2.0
		p1 = (prob.init_min[1] + prob.init_max[1]) / 2.0
		p2 = (prob.init_min[2] + prob.init_max[2]) / 2.0
		r = p0 - prob.init_min[0]
		for i in range(vector_input.shape[0]):
			x = vector_input[i, 0]
			y = vector_input[i, 1]
			z = vector_input[i, 2]
			cons = (x - p0) * (x - p0) + (y - p1) * (y - p1) + (z - p2) * (z - p2) <= r * r
			vector_input = vector_input[cons]

	print("vector_input:", vector_input)
	step_nums = 20
	color_base = np.linspace(0, vector_input.shape[0] - 1, vector_input.shape[0])
	x_pre = vector_input
	mlab.points3d(x_pre[:, 0], x_pre[:, 1], x_pre[:, 2], color_base, scale_mode="none", colormap="copper", scale_factor=0.25)
	for i in range(step_nums):
		x_next = prob.vector_field(x_pre)
		mlab.points3d(x_next[:, 0], x_next[:, 1], x_next[:, 2], color_base, scale_mode="none", colormap="copper", scale_factor=0.25)
		x_pre = x_next





	# x_data = np.linspace(prob.init_min[0], prob.init_max[0], 2)
	# y_data = np.linspace(prob.init_min[1], prob.init_max[1], 2)
	# z_data = np.linspace(prob.init_min[2], prob.init_max[2], 2)
	# s = np.array(np.meshgrid(x_data, y_data, z_data))
	# s = s.reshape(-1, 3)  # [[x,y,z], ..., [x,y,z]]
	# x = s[:, 0]
	# y = s[:, 1]
	# z = s[:, 2]
	#
	# # print(f"x:{x}, y:{y}, z:{z}")
	#
	# if prob.init_shape == 1:
	# 	x0 = torch.tensor(s, dtype=torch.float32).to("cpu")
	# if prob.init_shape == 2:
	# 	p0 = (prob.init_min[0] + prob.init_max[0]) / 2.0
	# 	p1 = (prob.init_min[1] + prob.init_max[1]) / 2.0
	# 	p2 = (prob.init_min[2] + prob.init_max[2]) / 2.0
	# 	r = p0 - prob.init_min[0]
	# 	cons = (x - p0) * (x - p0) + (y - p1) * (y - p1) + (z - p2) * (z - p2) <= r * r
	# 	x0 = torch.tensor(s[cons], dtype=torch.float32).to("cpu")
	# print(x0.shape[0])
	# color_base = np.linspace(0, x0.shape[0]-1, x0.shape[0])  # [1, 2, ..., 10]
	# # print(color_base)
	#
	# x_pre = x0
	# step_nums = 20
	# result = [x_pre]
	# color_each = color_base
	# for i in range(step_nums):
	# 	x_next = prob.vector_field(x_pre)
	# 	result.append(x_next)
	# 	color_each = np.concatenate((color_each, color_base), axis=0)  # [1, 2, ..., 10, ..., 1, 2, ..., 10]
	# 	x_pre = x_next
	#
	# result = torch.cat(result, 0).cpu().detach().numpy()
	# # print(color_each)
	# # print(result)
	#
	# # for i in range(result.shape[0]):
	# # 	print(i)
	# 	# print(f"route[0]: {route[:, 0]}, route[1]: {route[:, 1]}, route[2]: {route[:, 2]}")
	# mlab.points3d(result[:, 0], result[:, 1], result[:, 2], color_each, scale_mode="none", colormap="copper", scale_factor=0.25)







def mlab_plt_flow():
	vector_input = gen_plot_data()
	vector_field = prob.vector_field(vector_input)
	# attention! for plot-3d to work, prob.PLOT_LEN_V should be equal to prob.PLOT_LEN_B
	u = (vector_field[:, 0]).reshape(prob.PLOT_LEN_B[0], prob.PLOT_LEN_B[1], prob.PLOT_LEN_B[2])
	v = (vector_field[:, 1]).reshape(prob.PLOT_LEN_B[0], prob.PLOT_LEN_B[1], prob.PLOT_LEN_B[2])
	w = (vector_field[:, 2]).reshape(prob.PLOT_LEN_B[0], prob.PLOT_LEN_B[1], prob.PLOT_LEN_B[2])

	obj = flow(u, v, w, linetype='ribbon', seed_scale=1.0)
	obj.stream_tracer.maximum_propagation = 2000
	return obj


# plot frame
def mlab_plt_frame():
	x, y, z = np.mgrid[ (prob.domain_min[0]) : (prob.domain_max[0]) : complex(0, prob.PLOT_LEN_B[0]), \
							(prob.domain_min[1]) : (prob.domain_max[1]) : complex(0, prob.PLOT_LEN_B[1]), \
								(prob.domain_min[2]) : (prob.domain_max[2]) : complex(0, prob.PLOT_LEN_B[2]) ]
	u = 0 * x
	v = 0 * y
	w = 0 * z
	src = mlab.pipeline.vector_field(u, v, w)

	frame = mlab.pipeline.vectors(src, scale_factor=0, opacity=0)
	# frame = mlab.quiver3d(u, v, w, color = (1, 1, 1), scale_factor = 0, opacity = 0.0)
	# print(src.get_output_dataset())
	# input
	mlab.outline()
	# mlab.axes()


# plot 3d-system
def plot_barrier_3d(model):
	mlab.figure(fgcolor = (0, 0, 0), bgcolor = (1, 1, 1))
	mlab_plt_barrier(model)
	# mlab_plt_barrier_withoutModel()
	mlab_plt_init()
	mlab_plt_unsafe()
	mlab_plt_vector()
	# mlab_plt_flow()

	# mlab_plt_discrete_trajectory()
	mlab_plt_frame()

	# module 'mayavi.mlab' has no attribute 'xlim'
	# mlab.xlim(prob.domain_min[0], prob.domain_max[0])
	# mlab.ylim(prob.domain_min[1], prob.domain_max[1])
	# mlab.zlim(prob.domain_min[2], prob.domain_max[2])

	mlab.show()



# ###########################################
# # plot heart
# ###########################################
# def mlab_plt_heart():
# 	x, y, z = np.ogrid[ (prob.domain_min[0]) : (prob.domain_max[0]) : complex(0, prob.PLOT_LEN_B[0]), \
# 								(prob.domain_min[1]) : (prob.domain_max[1]) : complex(0, prob.PLOT_LEN_B[1]), \
# 									(prob.domain_min[2]) : (prob.domain_max[2]) : complex(0, prob.PLOT_LEN_B[2]) ]

# 	heart_scalar = (x*x+9/4*y*y+z*z-1) * (x*x+9/4*y*y+z*z-1) * (x*x+9/4*y*y+z*z-1) - x*x*z*z*z -9/80*y*y*z*z*z
# 	heart = mlab.contour3d(heart_scalar, contours = [0], color = (1,0,0), opacity = 0.3)

if __name__ == "__main__":
    model = cdinn2.gen_nn()
    model.load_state_dict(torch.load('./model/cdinn_c9_1_15_0.1_0.05_1_train.pt'), strict=True)
	# express = (-0.54543733999999999 + math.sin(math.sin(((x0 - math.tanh((0.82548980000000005 + x1))) / (1.1922542933451701 + 1.33552912280763 * math.exp(( - x1 + x2)))))))
    plot_barrier_3d(model)
    # plot_barrier_3d()
	# plot_barrier_3d(express)
    pass
